import os
from CalTorque import CalTorque
from datetime import datetime
import matplotlib
from matplotlib import pyplot as PLT
from scipy.optimize import curve_fit
from scipy.stats import linregress
from datetime import datetime
data_files = []
dp = 'data/'
ls= os.listdir(dp)
for i in ls:
   if i.split('.')[-1] not in ['xlsx', 'xls']: continue
   if 'out' in i: continue
   elif 'in' not in i: continue
   data_files.append(dp+i)

max_torque = []
max_at_B2 = []
max_at_B3 = []
date = []
intdate = []
filelist = []
linfit = []
linfit_B2 = []
linfit_B3 = []

for f in data_files:   
    data = CalTorque(f)
    if data['source'] != 'weakTh-228': continue
    if data['station'] != 'S5': continue
    if max(data['pos_offsetfix']) < 350: continue   # because some stations in run_history are wrong
    if max(data['pos_offsetfix']) > 400: continue
    print f
    
    filelist.append(f.split('.')[0].split(' ')[2].split('_')[0])
    max_torque.append(data['max_torque_inoz'])
    rundate = datetime.strptime(str(data['datetime'][0]).split(' ')[0], '%Y-%m-%d')
    date.append(rundate)
    
    pos_torque = []
    for i, x in enumerate(data['pos_offsetfix']):
        if x > 300 and x < 330:
            pos_torque.append(data['torque_inoz'][i])
    max_at_B2.append(max(pos_torque))
    
    pos_torque2 = []
    for i, x in enumerate(data['pos_offsetfix']):
        if x > 350 and x < 375:
            pos_torque2.append(data['torque_inoz'][i])
    max_at_B3.append(max(pos_torque2))



for d in date:
    intd = d - min(date)
    intdate.append(intd.days)

lin = linregress(intdate, max_torque)
lin_B2 = linregress(intdate, max_at_B2)
lin_B3 = linregress(intdate, max_at_B3)


def Polyfnc(x, m, b):
    return m*x + b

for thing in intdate:
    linfit.append(Polyfnc(thing, lin[0], lin[1]))
    linfit_B2.append(Polyfnc(thing, lin_B2[0], lin_B2[1]))
    linfit_B3.append(Polyfnc(thing, lin_B3[0], lin_B3[1]))

fig = PLT.figure(figsize=(15,8),dpi=150)
box = [0.14, 0.14, 0.76, 0.76]
ax1 = fig.add_axes(box)
ax1.set_ylabel('Max Torque (in.oz)')
ax1.set_xlabel('Date')
PLT.title('Max Torque Versus Date')
ax1.grid()
ax1.scatter(date, max_torque, color = 'g', label = 'Data')
#for i, txt in enumerate(filelist):
    #ax1.annotate(txt, (date[i], max_torque[i]), position = (date[i], max_torque[i]), size = 'x-small', alpha = 0.5, rotation = 45)
ax1.plot(date, linfit, '-', color = 'g', label = 'Linear fit: y = %.2f * x + %.2f' % (lin[0], lin[1]))
h, l = ax1.get_legend_handles_labels()
PLT.legend(h, l, 'upper left')
fig.savefig('maxt_vs_date.png')

fig2 = PLT.figure(figsize=(15,8),dpi=150)
box = [0.14, 0.14, 0.76, 0.76]
ax2 = fig2.add_axes(box)
ax2.set_ylabel('Max Torque (in.oz)')
ax2.set_xlabel('Date')
PLT.title('Max Torque Versus Date for the Second Inner Bend')
ax2.grid()
ax2.scatter(date, max_at_B2, color = 'b', label = 'Data')
ax2.plot(date, linfit_B2, '-', color = 'b', label = 'Linear fit: y = %.2f * x + %.2f' % (lin_B2[0], lin_B2[1]))
h, l = ax2.get_legend_handles_labels()
PLT.legend(h, l, 'upper left')
fig2.savefig('maxt_vs_date_for_B2.png')

fig3 = PLT.figure(figsize=(15,8),dpi=150)
box = [0.14, 0.14, 0.76, 0.76]
ax3 = fig3.add_axes(box)
ax3.set_ylabel('Max Torque (in.oz)')
ax3.set_xlabel('Date')
PLT.title('Max Torque Versus Date for the Third Inner Bend')
ax3.grid()
ax3.scatter(date, max_at_B3, color = 'r', marker = 'D', label = 'Data')
ax3.plot(date, linfit_B3, '-', color = 'r', label = 'Linear fit: y = %.2f * x + %.2f' % (lin_B3[0], lin_B3[1]))
h, l = ax3.get_legend_handles_labels()
PLT.legend(h, l, 'upper left')
fig3.savefig('maxt_vs_date_for_B3.png')

fig4 = PLT.figure(figsize=(15,8),dpi=150)
box = [0.14, 0.14, 0.76, 0.76]
ax4 = fig4.add_axes(box)
ax4.set_ylabel('Max Torque (in.oz)')
ax4.set_xlabel('Date')
PLT.title('Max Torque Versus Date for Bends A6 and A7')
ax4.grid()
ax4.scatter(date, max_at_B3, color = 'r', marker = 'D', label = 'Bend A7')
ax4.plot(date, linfit_B3, '-', color = 'r', label = 'Bend A7 linear fit: y = %.2f * x + %.2f' % (lin_B3[0], lin_B3[1]))
ax4.scatter(date, max_at_B2, color = 'b', label = 'Bend A6')
ax4.plot(date, linfit_B2, '-', color = 'b', label = 'Bend A6 linear fit: y = %.2f * x + %.2f' % (lin_B2[0], lin_B2[1]))
h, l = ax4.get_legend_handles_labels()
PLT.legend(h, l, 'upper left')
fig4.savefig('maxt_vs_date_for_peaks.png')
